"""
题目：给定一个图，根据其是否为有向图，统计边的总数（无向图中A-B算 1 条边，有向图中A→B和B→A算 2 条边）。
"""
class Graph:
    def __init__(self, is_directed=False):
        self.vertices = {}
        self.is_directed = is_directed

    def add_vertex(self, vertex):
        if vertex not in self.vertices:
            self.vertices[vertex] = []

    def add_edge(self, v1, v2):
        self.add_vertex(v1)
        self.add_vertex(v2)
        self.vertices[v1].append(v2)
        if not self.is_directed:
            self.vertices[v2].append(v1)

    def count_total_edges(self):
        """统计图中边的总数"""
        edge_count = 0
        if self.is_directed:
            # 有向图：直接累加所有邻接列表长度（每条边仅在起点邻接表中记录）
            for vertex in self.vertices:
                edge_count += len(self.vertices[vertex])
        else:
            # 无向图：累加后除以2（避免重复计数，如A-B在A和B的邻接表中各记1次）
            for vertex in self.vertices:
                edge_count += len(self.vertices[vertex])
            edge_count //= 2
        return edge_count

# 测试
if __name__ == "__main__":
    # 测试无向图
    undir_graph = Graph()
    undir_graph.add_edge(0, 1)
    undir_graph.add_edge(0, 2)
    undir_graph.add_edge(1, 2)
    print("无向图边总数:", undir_graph.count_total_edges())  # 输出: 3

    # 测试有向图
    dir_graph = Graph(is_directed=True)
    dir_graph.add_edge('A', 'B')
    dir_graph.add_edge('B', 'A')
    dir_graph.add_edge('B', 'C')
    print("有向图边总数:", dir_graph.count_total_edges())  # 输出: 3（A→B、B→A、B→C）